/*
  Copyright (c) 2010 Davide Rossetti (davide.rossetti@roma1.infn.it)

  This file is part of CuHB (CUDA Heisenberg) package.
  
  CuHB is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  CuHB is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
  
  You should have received a copy of the GNU General Public License
  along with CuHB.  If not, see <http://www.gnu.org/licenses/>.

*/

#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <string.h>

#include <cuda_runtime.h>
#include <math.h>

#include "hb.h"
#include "rng.hpp"

//-----------------------------------------------------------------------------
//-----------------------------------------------------------------------------

Jmat::Jmat(Config &cfg, int pad)
{
    assert(cfg.dims == 3);
    assert(cfg.cpu_precision == cfg.gpu_precision);
    assert(cfg.cpu_precision == SinglePrecision);

    m_cfg = cfg;

    m_volume = 1;
    for(int i=0; i<m_cfg.dims; ++i) {
        m_volume *= m_cfg.lat[i];
    }
    m_pad = pad;
    m_stride = m_volume + m_pad;

    // float[2] J[NDIMS,VOL]
    size_t gpu_matrix_size = 2 * sizeof(float4);
    m_gpu_bytes = m_stride * m_cfg.dims * gpu_matrix_size;

    // float[3][3] J[NDIMS,VOL]
    //size_t cpu_matrix_element_size = sizeof(float);
    //size_t cpu_matrix_size = m_cfg.dims * m_cfg.dims * cpu_matrix_element_size;
    //m_cpu_bytes = m_volume * m_cfg.dims * cpu_matrix_size;
    m_cpu_bytes = m_gpu_bytes;

    m_cpu_buf = malloc(m_cpu_bytes);
    if(m_cpu_buf == NULL) {
        error("Jmat: error allocating CPU memory\n");
    }
    memset(m_cpu_buf, 0, m_cpu_bytes);

    // zero buf
    if(cudaMalloc((void**)&m_gpu_buf, m_gpu_bytes) != cudaSuccess) {
        error("Svec: error allocating GPU memory\n");
    }
    if(cudaMemset(m_gpu_buf, 0, m_gpu_bytes) != cudaSuccess) {
        error("Svec: error setting GPU memory\n");
    }
}

//-----------------------------------------------------------------------------

Jmat::~Jmat()
{
    cudaFree(m_gpu_buf);
    m_gpu_buf = NULL;

    free(m_cpu_buf);
    m_cpu_buf = NULL;
}

//-----------------------------------------------------------------------------

void Jmat::init_ident()
{
    assert(m_cpu_buf);
    int lx = m_cfg.lat[DIR_X];
    int ly = m_cfg.lat[DIR_Y];
    int lz = m_cfg.lat[DIR_Z];
    float4* vec = (float4*)m_cpu_buf;

    trace("packing 3x3 Identity into 2 x float4\n");

    /*
      2 x float4
      only use first 3 floats of each float4
      to encode above diagonal part

      1.0 0.0 0.0
      Ident = 0.0 1.0 0.0
      0.0 0.0 1.0
    */
    float4 f1 = make_float4(1.0, 0.0, 0.0, 0.0);
    float4 f2 = make_float4(1.0, 0.0, 1.0, 0.0);

    for(int ix=0; ix<lx; ++ix)
        for(int iy=0; iy<ly; ++iy)
            for(int iz=0; iz<lz; ++iz) {
                int idx = (iz) + (iy)*lz + (ix)*(lz*ly);
                for(int dir = DIR_X; dir < 3; ++dir) {
                    vec[idx + ((dir*2+0)*m_stride)] = f1;
                    vec[idx + ((dir*2+1)*m_stride)] = f2;
                    assert(sizeof(float4)*(idx + ((dir*2+1)*m_stride)) < m_gpu_bytes);
                }
            }
}

//-----------------------------------------------------------------------------

void Jmat::init_gaussian(rng_uniform &rng)
{
    assert(m_cpu_buf);
    int lx = m_cfg.lat[DIR_X];
    int ly = m_cfg.lat[DIR_Y];
    int lz = m_cfg.lat[DIR_Z];
    float4* vec = (float4*)m_cpu_buf;

    trace("packing Gaussian into 2 x float4, stride=%d\n", m_stride);

    for(int ix=0; ix<lx; ++ix)
        for(int iy=0; iy<ly; ++iy)
            for(int iz=0; iz<lz; ++iz) {
                int idx = (iz) + (iy)*lz + (ix)*(lz*ly);
                for(int dir = DIR_X; dir < 3; ++dir) {
                    float4 f1 = make_float4(rng.gen_gaussian_float(), 
                                            rng.gen_gaussian_float(), 
                                            rng.gen_gaussian_float(), 
                                            0.0);
                    float4 f2 = make_float4(rng.gen_gaussian_float(), 
                                            rng.gen_gaussian_float(), 
                                            rng.gen_gaussian_float(), 
                                            0.0);
                    vec[idx + ((dir*2+0)*m_stride)] = f1;
                    vec[idx + ((dir*2+1)*m_stride)] = f2;
                    assert(sizeof(float4)*(idx + ((dir*2+1)*m_stride)) < m_gpu_bytes);
                }
            }
}

//-----------------------------------------------------------------------------

void Jmat::cpu_to_gpu()
{
    trace("copying cpu %p -> GPU %p  %d bytes\n", m_cpu_buf, m_gpu_buf, m_gpu_bytes);
    if(cudaMemcpy(m_gpu_buf, m_cpu_buf, m_gpu_bytes, cudaMemcpyHostToDevice) != cudaSuccess) {
        error("%s: error copying\n", __FUNCTION__);
	}
}

//-----------------------------------------------------------------------------

void Jmat::gpu_to_cpu()
{
    if(cudaMemcpy(m_cpu_buf, m_gpu_buf, m_cpu_bytes, cudaMemcpyDeviceToHost) != cudaSuccess) {
        error("%s: error copying\n", __FUNCTION__);
	}
}

//-----------------------------------------------------------------------------
//-----------------------------------------------------------------------------

/*
 * Local variables:
 *  mode: c++
 *  c-indent-level: 4
 *  c-basic-offset: 4
 *  tab-width: 4
 *  indent-tabs-mode: nil
 * End:
 */
